YOLO v8 自訂資料模型訓練實作

yolo python ai

I code so I am 2023-02-22 07:33:37 ‧ 11431 瀏覽

分享至

前言

前一篇說明 YOLO v8 的安裝與預測，預設的模型可辨識COCO資料集的80種物件類別，如果要辨識其他物件類別，就必須準備自訂資料集(Custom dataset)，訓練專屬的模型。

準備自訂資料集

Roboflow提供許多公用資料集，適用各種類別的辨識，涵蓋物件偵測(object detection)、分類(classification)、實例分割(instance segmentation)及語義分割(semantic segmentation)，這裡採用『美元紙鈔』(Dollar bill)資料集，包括1、5、10、20、50元的正反面的圖片，共10個類別，部分圖片含多張紙鈔。

自這裡下載後，解壓縮至一目錄，假設為 Dollar_Bill，需先進行以下前置處理：

在Dollar_Bill目錄建一子目錄datasets，並將train、valid、test目錄移放在datasets下。
修改data.yaml，將前3行的『..』改為『.』。

模型訓練

在Dollar_Bill目錄執行下列指令，可進行模型訓練：

yolo task=detect mode=train model=yolov8s.pt data=./data.yaml epochs=5 imgsz=640

筆者PC配備GTX1050Ti GPU，每一執行週期(Epoch)約訓練3分鐘，5個執行週期約15分鐘，準確率(mAP)還不錯。

訓練成功後，最佳模型會存為 Dollar_Bill\runs\detect\train\weights\best.pt，另外還有一個最新的模型last.pt。

模型評分

在Dollar_Bill目錄執行下列指令，可進行模型訓練：

yolo task=detect mode=val model=./runs/detect/train/weights/best.pt data=./data.yaml

結果如下，準確率(mAP)還不錯：

偵測圖片的結果存在 Dollar_Bill/runs/detect/val目錄，可觀察每一個jpg檔案，大部分都可把紙鈔加框，結果如下：

預測

在Dollar_Bill目錄執行下列指令，可進行模型預測，以下指令預測./datasets/test/images目錄下所有檔案：

yolo task=detect mode=predict model=./runs/detect/train/weights/best.pt conf=0.25 source=./datasets/test/images

結果如下：

也可以使用前一篇的程式：

from ultralytics import YOLO
from PIL import Image
import cv2

model = YOLO("./runs/detect/train/weights/best.pt")

# from PIL
im1 = Image.open(r".\Dollar_Bill\datasets\test\images\IMG_1909_jpg.rf.3f2bde8dc899b141a4361677e55242e4.jpg")
# save=True：存檔
results = model.predict(source=im1, save=True)

結果會存檔位置預設在 runs\detect\predict 資料夾下。

結語

整體而言，YOLO v8 的測試比前面版本方便，訓練時間也縮短很多，準確率也不錯，無特殊異常狀況需處理。
除了物件偵測外，YOLO v8 還提供分類(Classification)、語義分割(Semantic segmentation)及物件追蹤(Object tracking)功能，可參閱官方文件，下次再試試看。

以下為工商廣告:)。
深度學習PyTorch入門到實戰應用影音課程：